/*
 * @Author: your name
 * @Date: 2020-11-04 03:29:27
 * @LastEditTime: 2021-11-11 14:11:58
 * @LastEditors: Yang jin
 * @Description: In User Settings Edit
 * @FilePath: /fast-fasta2kmers/src/fasta2kmers.cpp
 */
#include <errno.h>
#include <stdio.h>
#include <zlib.h>
#include <time.h>
#include <random>

#include <dirent.h>
#include <unistd.h>
#include <sys/types.h>
#include <ctype.h>

#include <stdlib.h>
#include <time.h>

#include <vector>
#include <thread>

// #include<getopt.h>
// extern int optind,opterr,optopt;
// extern char *optargi;
//定义长选项
// static struct option long_options[] =
// {
//     {"help",no_argument,NULL,'u'},
//     {"mode",required_argument,NULL,'m'},
//     {"k",required_argument,NULL,'k'},
//     {"k_low",required_argument,NULL,'l'},
//     {"k_high",required_argument,NULL,'h'},
//     {"inpath",required_argument,NULL,'i'},
//     {"outpath",required_argument,NULL,'o'}
// };

#include <sys/stat.h>
#include "kseq.h"

KSEQ_INIT(gzFile, gzread)

char *regchar = "AGCTN";
char *newchar = "N";
bool regchar_mp[128] = {false};

void usage()
{
    fprintf(stderr, "Usage: ./fast-fasta2kmers <in.fa.gz> <out.fa> <mode> <k>(mode:0 non-overlapping kmers;1 overlapping kmers;2：nonoverlapping kmer; 3: variable overlapping kmers;4:overlapping pair kmers;5:non-overlap pair kmers)\n");

}

bool is_valid(int argc)
{
    if (argc > 6)
    {
        usage();
        return 0;
    }
    return 1;
}

int strtrans(char *str, char *regchar, char *newchar)
{
    if (str == NULL)
    {
        return 0;
    }
    int cnt =0;
    for (int i = 0; str[i] != 0; i++)
    {
        str[i] = toupper(str[i]);
        if (!regchar_mp[str[i]])
        {
            str[i] = 'N';
            cnt++;
        }
        if (str[i]=='N')
        {
            cnt++;
        }
    }
    return cnt;
}

void read_length(FILE *f,char *name,int l){
    fprintf(f, ">%s\n", name);
    fprintf(f,"%d\n",l);
    fflush(f);
}

void split_pair_non_kmer_seq(FILE *f, char *name, char *seq, int l, int k)
{
    char buffer[k + 1];
    buffer[k] = 0;
    int mid = (l ) / 2;

    if (int(l / 2) < k)
    {
        return;
    }
    int cnt =strtrans(seq,regchar,newchar);
    if (cnt > l/4)
    {
        return;
    }

    fprintf(f, ">%s\n", name);

    int i= 0;
    for (; i < mid; i+=k)
    {
        memcpy(buffer, seq + i, k * sizeof(char));
        fprintf(f, "%s ", buffer);
    }
    fprintf(f, "\t");
    for (; i < l ; i+=k)
    {
        memcpy(buffer, seq + i, k * sizeof(char));
        fprintf(f, "%s ", buffer);
    }

    fprintf(f, "\n");

    fflush(f);
}

void split_pair_seq(FILE *f, char *name, char *seq, int l, int k)
{
    char buffer[k + 1];
    buffer[k] = 0;
    int mid = (l - k + 1) / 2;

    if (int(l / 2) < k)
    {
        return;
    }
    int cnt =strtrans(seq,regchar,newchar);
    if (cnt > l/4)
    {
        return;
    }

    fprintf(f, ">%s\n", name);

    for (int i = 0; i < mid; i++)
    {
        memcpy(buffer, seq + i, k * sizeof(char));
        fprintf(f, "%s ", buffer);
    }

    fprintf(f, "\t");

    for (int i = mid; i < l - k + 1; i++)
    {
        memcpy(buffer, seq + i, k * sizeof(char));
        fprintf(f, "%s ", buffer);
    }

    fprintf(f, "\n");

    fflush(f);
}

void split_stride_seq(FILE *f, char *name, char *seq, int l,int stride, int k)
{
    // int cnt =strtrans(seq,regchar,newchar);
    // if (cnt > l/2)
    // {
    //     return;
    // }
    
    char buffer[k + 1];
    buffer[k] = 0;
    int index = 0;

    for (int i = 0; i < l; i+=stride)
    {
        if ((l-i)<k/2){
            break;
        }
        fprintf(f, ">index|%d|%s\n",index,name);
        index ++;
        memcpy(buffer, seq + i, k * sizeof(char));
        fprintf(f, "%s\n", buffer);
    }


    fflush(f);
}

void split_seq(FILE *f, char *name, char *seq, int l, int k)
{
    // int cnt =strtrans(seq,regchar,newchar);
    // if (cnt > l/2)
    // {
    //     return;
    // }
    
    char buffer[k + 1];
    buffer[k] = 0;
    fprintf(f, ">%s\n", name);

    for (int i = 0; i < l - k + 1; i++)
    {
        memcpy(buffer, seq + i, k * sizeof(char));
        fprintf(f, "%s ", buffer);
    }

    fprintf(f, "\n");

    fflush(f);
}

void split_non_overlap_kmers(FILE *f,char *name,char *seq,int l,int k)
{
     int cnt =strtrans(seq,regchar,newchar);
    if (cnt > l/10)
    {
        return;
    }
    
    char buffer[k + 1];
    buffer[k] = 0;
    fprintf(f, ">%s\n", name);

    for (int i = 0; i < l; i+=k)
    {
        memcpy(buffer, seq + i, k * sizeof(char));
        fprintf(f, "%s ", buffer);
    }

    fprintf(f, "\n");
    fflush(f);

}

void split_overlap_seq_pair(FILE *f, char *name, char *seq1, char *seq2, int l1, int l2, int k)
{
    char buffer[k + 1];
    buffer[k] = 0;
    fprintf(f, ">%s\n", name);
    for (int i = 0; i < l1 - k + 1; i++)
    {
        memcpy(buffer, seq1 + i, k * sizeof(char));
        fprintf(f, "%s ", buffer);
    }
    fprintf(f, "\t");
    for (int i = 0; i < l2 - k + 1; i++)
    {
        memcpy(buffer, seq2 + i, k * sizeof(char));
        fprintf(f, "%s ", buffer);
    }
    fprintf(f, "\n");
    fflush(f);
}

void split_non_overlap_seq(FILE *f, char *name, char *seq, int l, int k, int minlength)
{
    if (l < minlength)
        return;

    //将非AGCTN的词汇转换为N（未知词汇）,如果未知词汇量占据1/4以上，就直接舍弃
    int cnt =strtrans(seq,regchar,newchar);
    if (cnt > l/2)
    {
        return;
    }
    
    char buffer[k + 1];
    buffer[k] = 0;
    int i;
    int index=0;
    for (i = 0; i < l - k; i += k)
    {
        fprintf(f, ">index|%d|%s\n",index,name);
        index ++;
        memcpy(buffer, seq + i, k * sizeof(char));
        fprintf(f, "%s\n", buffer);
    }
    //discard the remained sequence whose length is lower than k/2

    buffer[l - i] = 0;
    fprintf(f, ">index|%d|%s\n", index,name);
    memcpy(buffer, seq + i, (l - i) * sizeof(char));
    fprintf(f, "%s\n", buffer);
    // }

    fflush(f);
}

void split_variable_overlap_seq(FILE *f, char *name, char *seq, int l, int k_low, int k_high)
{
    /** Split the sequence into variable length mers, k is sampled from the discrete uniform distribution int Uniform(k_low,k_high)
     *  f(s) = (S_1~1+k1,S_2~2+k2,,,,S_n-k_high+1 ~ n-k_high+1+k_n)   
     */

    std::random_device rd;  //Will be used to obtain a seed for the random number engine
    std::mt19937 gen(rd()); //Standard mersenne_twister_engine seeded with rd()
    std::uniform_int_distribution<> distrib(k_low, k_high);
    char buffer[k_high + 1];
    buffer[k_high] = 0;
    fprintf(f, ">%s\n", name);
    int k;
    for (int i = 0; i < l - k_high; i++)
    {
        k = distrib(gen);
        buffer[k] = 0;
        memcpy(buffer, seq + i, k * sizeof(char));
        fprintf(f, "%s ", buffer);
    }
    fprintf(f, "\n");
    fflush(f);
}

void get_kmers(char *argv[], FILE *f, char *name, char *seq, int l, int minlength = 10, char *seq2 = "", int l2 = 0)
{
    int mode = atoi(argv[3]);
    int k, k_low, k_high,stride;
    switch (mode)
    {
    case 0:
        k = atoi(argv[4]);
        split_non_overlap_seq(f, name, seq, l, k, minlength);
        break;
    case 1:
        k = atoi(argv[4]);
        split_seq(f, name, seq, l, k);
        break;
    case 2:
        k = atoi(argv[4]);
        split_non_overlap_kmers(f,name,seq,l,k);
        break;
    case 3:
        k_low = atoi(argv[4]);
        k_high = atoi(argv[5]);
        if (k_low >= k_high)
        {
            fprintf(stderr, "k_low  should be less than k_high ,Do you mean change the k_low = %d k_high= %d? \n", k_high, k_low);
            int tmp = k_low;
            k_low = k_high;
            k_high = tmp;
        }

        split_variable_overlap_seq(f, name, seq, l, k_low, k_high);
        break;

    case 4:
        k = atoi(argv[4]);
        split_pair_seq(f, name, seq, l, k);
        break;
    
    case 5:
        k = atoi(argv[4]);
        split_pair_non_kmer_seq(f, name, seq, l, k);
        break;
    
    case 6:
        read_length(f,name,l);
        break;

    case 7:
        stride = atoi(argv[5]);
        k = atoi(argv[4]);
        split_stride_seq(f,name,seq,l,stride,k);
        break;

    default:
        usage();
        break;
    }
}

int main(int argc, char *argv[])
{
    if (!is_valid(argc))
        return 1;

    regchar_mp['A'] = true;
    regchar_mp['G'] = true;
    regchar_mp['C'] = true;
    regchar_mp['T'] = true;
    regchar_mp['N'] = true;

    gzFile fp;
    FILE *fw;
    kseq_t *seq;
    int l;

    double begintime, endtime;
    begintime = clock();

    char const *inpath = argv[1];
    struct stat inpath_buf;
    stat(inpath, &inpath_buf);

    //The input path is file
    char *fn = argv[1];
    char *fo = argv[2];

    fp = gzopen(fn, "r");
    fw = fopen(fo, "w");

    if (fp == nullptr)
    {
        fprintf(stderr, "Couldn't open %s : %s\n", fn, errno ? strerror(errno) : "Out of memory");
        return 1;
    }

    seq = kseq_init(fp);

    while ((l = kseq_read(seq)) >= 0)
    {
        get_kmers(argv, fw, seq->name.s, seq->seq.s, l);
    }

    kseq_destroy(seq);
    gzclose(fp);
    fclose(fw);

    endtime = clock();

    fprintf(stdout, "Splited Done!\nCost total %.3f secs\n", (endtime - begintime) / CLOCKS_PER_SEC);

    return 0;
}